"""
对测试集和验证集进行2D裁剪
并将2D文件idd用txt文件保存下来，2D-cut后的数据放在一个文件夹下（跟训练集2D-cut的数据同一文件夹）
"""
import os
import sys
sys.path.append(os.path.split(sys.path[0])[0])
import SimpleITK as sitk
from multiprocessing.dummy import Pool
from config import zunYiParameter as para
import numpy as np
def process(file):
    # 将CT和标签入读内存
    print(file)
    volume = sitk.ReadImage(os.path.join(para.nii_volume_path, file), sitk.sitkInt16)
    volume_array = sitk.GetArrayFromImage(volume)
    #print('处理前的volume shape:',volume_array.shape,end = '   ')
    seg = sitk.ReadImage(os.path.join(para.nii_seg_path, file.replace('volume', 'segmentation')), sitk.sitkUInt8)
    seg_array = sitk.GetArrayFromImage(seg)
    seg_file_name=file.replace('volume', 'segmentation')
    ct_file_name=file
    # seg_array=processLiverOrTumorMask(para.Liver_or_Tumor,seg_array)
    # 将灰度值在阈值之外的截断掉
    # volume_array[volume_array > para.upper] = para.upper
    # volume_array[volume_array < para.lower] = para.lower

    all_indexs = [i for i in range(len(volume_array))]
    #保存图片
    file_names=[]

    for index in all_indexs:
        ct_name=os.path.join("volume",str(index)+"_"+ct_file_name)
        seg_name=os.path.join("segmentation", str(index) + "_" + seg_file_name)
        file_names.append(ct_name+" "+seg_name)
        #保存
        ct_path = os.path.join(para.cut2d_save_path,ct_name)
        seg_path = os.path.join(para.cut2d_save_path,seg_name)
        sitk.WriteImage(sitk.GetImageFromArray(volume_array[index]),ct_path)
        sitk.WriteImage(sitk.GetImageFromArray(seg_array[index]), seg_path)

    return file_names


def main():
    if not os.path.exists(para.cut2d_save_path_vol):
        os.makedirs(para.cut2d_save_path_vol)
    if not os.path.exists(para.cut2d_save_path_seg):
        os.makedirs(para.cut2d_save_path_seg)

    if not os.path.exists(os.path.dirname(para.test2d_id_path)):
        os.makedirs(os.path.dirname(para.test2d_id_path))

    with open(para.test_nii_id_path, 'r') as f:
        test_ids = f.read().splitlines()
    test_file_list=[os.path.basename(file.split()[0]) for file in test_ids]
    print("process nii file number:", len(test_file_list))
    #处理测试集
    pool_for_test = Pool(2)
    result=pool_for_test.map(process, test_file_list)
    test_file_names_flatten = []
    for i in range(len(result)):
        test_file_names_flatten += result[i]
    test_file_names_flatten = np.array(test_file_names_flatten)
    print("test size", test_file_names_flatten.shape)
    pool_for_test.close()
    pool_for_test.join()


    np.savetxt(para.test2d_id_path, test_file_names_flatten, fmt = "%s")

if __name__ == '__main__':
    '''
    test size (768,)
    '''
    main()